## Get VF / Policy function

# Specify here the type of every input
function get_VF_baseline(;param,guess_exp_VF,exp_profits,verbose=false,crit = 1e-6,max_iter = 500)
	
	# initialize objects 
	exp_VF_new = copy(guess_exp_VF) 
	exp_VF_old = copy(guess_exp_VF)

	iter = 1
	VF_diff = Inf 
	while VF_diff > crit && iter < max_iter

		# print progress
		if verbose
			println("Iteration ",iter," with difference: ",VF_diff)
		end 
		
		# get exit probability given new guess 
		exit_proba = 1.0 .- exp.(.-exp.(.-(beta.*exp_VF_new .- param.level_fcost)./param.scale_fcost))
		exp_fcost = compute_exp_fcost(exit_proba = exit_proba, exp_VF = exp_VF_new, param = param)

		# compute new VF (over all z)
		VF_new = exp_profits .+ (1.0 .- exit_proba).*( .-exp_fcost .+ beta.*exp_VF_new )
			
		# save old results 
		exp_VF_old = copy(exp_VF_new)
			
		# compute new expected VF 
		exp_VF_new = transition_z * VF_new 
			
		# Update critical value + iteration
		VF_diff = maximum(abs.(exp_VF_new - exp_VF_old))
		iter = iter + 1		
	end

	## Check if cost parameters need to be indexed by wage!! 
  
	# update results
	exit_proba = 1.0 .- exp.(.-exp.(.-(beta.*exp_VF_new .- param.level_fcost)./param.scale_fcost))
	exp_fcost = compute_exp_fcost(exit_proba = exit_proba, exp_VF = exp_VF_new, param = param)
	VF = exp_profits .+ (1.0 .- exit_proba).*( .-exp_fcost .+ beta.*exp_VF_new )
  
  	# Return results 
	return (exp_VF = exp_VF_new, exit_proba = exit_proba, VF = VF, exp_fcost = exp_fcost)
end

# Specify here the type of every input
function get_VF_cf(;w,Y,type = "NC",tax=nothing,subsidy=nothing,param,guess_exp_VF,verbose=false,crit = 1e-6,max_iter = 500, n_eps = 500)
	
	# initialize objects 
	exp_VF_new = copy(guess_exp_VF)
	exp_VF_old = copy(guess_exp_VF)

	if tax == nothing
        tax = copy(vat_tax) 
    end 

	# compute z grid based on baseline results & choice of sigma 
	x_bar_baseline = (baseline_aggregates_goods.total_output)^(1/param.sigma)
	z_grid = (z_star_grid ./ x_bar_baseline).^(param.sigma/(param.sigma-1))

	# compute z_grid as function of sigma (to be able to vary in counterfactuals) 

	# compute profits
	if type == "NC"
		profit_results = compute_profits_NC(w = w, Y = Y, tax = tax, param = param, z_grid = z_grid)  
		# save results objects 
		exp_profits = profit_results.profits 
		revenue = profit_results.revenue 
		output_aggr = profit_results.output_aggr

	elseif type == "C" 
		profit_results = compute_profits_grid_cf(w = w, Y = Y, tax = tax, n_eps = n_eps, param = param, z_grid = z_grid)
		# save results objects 
		exp_profits = profit_results.expected_profits
		revenue = profit_results.expected_revenue
		output_aggr = profit_results.expected_revenue_output

	elseif type == "C-fix"
		if subsidy == nothing 
			println("Error: Need to specify subsidy rate")
		end 
		profit_results = compute_profits_grid_cf_fix_subsidy(w = w, Y = Y, subsidy = subsidy, tax = nothing, param = param, z_grid = z_grid)
		# save results objects 
		exp_profits = profit_results.expected_profits
		revenue = profit_results.expected_revenue
		output_aggr = profit_results.expected_revenue_output

	else 
		println("No other options than type == {C,NC}")
	end

	iter = 1
	VF_diff = Inf 
	while VF_diff > crit && iter < max_iter

		# print progress
		if verbose
			println("Iteration ",iter," with difference: ",VF_diff)
		end 
		
		# get exit probability given new guess 
		exit_proba = 1.0 .- exp.(.-exp.(.-(beta.*exp_VF_new .- param.level_fcost)./param.scale_fcost))
		exp_fcost = compute_exp_fcost(exit_proba = exit_proba, exp_VF = exp_VF_new, param = param)

		# compute new VF (over all z)
		VF_new = exp_profits .+ (1.0 .- exit_proba).*( .-exp_fcost .+ beta.*exp_VF_new )
			
		# save old results 
		exp_VF_old = copy(exp_VF_new)
			
		# compute new expected VF 
		exp_VF_new = transition_z * VF_new 
			
		# Update critical value + iteration
		VF_diff = maximum(abs.(exp_VF_new - exp_VF_old))
		iter = iter + 1		
	end

	## Check if cost parameters need to be indexed by wage!! 
  
	# update results
	exit_proba = 1.0 .- exp.(.-exp.(.-(beta.*exp_VF_new .- param.level_fcost)./param.scale_fcost))
	exp_fcost = compute_exp_fcost(exit_proba = exit_proba, exp_VF = exp_VF_new, param = param)
	VF = exp_profits .+ (1.0 .- exit_proba).*( .-exp_fcost .+ beta.*exp_VF_new )
  
  	# Return results 
	return (exp_VF = exp_VF_new, exit_proba = exit_proba, VF = VF, exp_fcost = exp_fcost, 
			exp_profits = exp_profits, revenue = revenue, output_aggr = output_aggr)
end

# find VF along transition 
function get_VF_transition(;w_path,Y_path, param, VF_results_end = nothing, profits_object_end = nothing, guess_exp_VF = nothing, verbose=false)

	### Step 0: Initialize objects and functions 

	# get length of transition 
	length_transition = length(w_path)

	# compute z grid based on baseline results & choice of sigma 
	x_bar_baseline = (baseline_aggregates_goods.total_output)^(1/param.sigma)
	z_grid = (z_star_grid ./ x_bar_baseline).^(param.sigma/(param.sigma-1))

	# get yearly VF (and other objects) 
	function get_yearly_VF(; w::Float64, Y::Float64, next_VF_object)

		# Step 1: compute profits 
		within_period_profits = compute_profits_NC(w = w, Y = Y, param = param, z_grid = z_grid) 
		profits = within_period_profits.profits 

		# Step 2: compute new value 
		VF_new = profits .+ (1.0 .- next_VF_object.exit_proba).*( .- next_VF_object.exp_fcost .+ (beta .* next_VF_object.exp_VF))

		# Step 3: construct new objects
		exp_VF_new = transition_z * VF_new 
		new_exit_proba = 1.0 .- exp.(.-exp.(.-(beta.*exp_VF_new .- param.level_fcost)./param.scale_fcost))
		exp_fcost = compute_exp_fcost(exit_proba = new_exit_proba, exp_VF = exp_VF_new, param = param)
		
		# Return results 
		return (VF = VF_new, exit_proba = new_exit_proba, exp_VF = exp_VF_new, 
				exp_fcost = exp_fcost, within_period_objects = within_period_profits) 
	end 

	# compute (stationary) VF at the end if not provided 
	if VF_results_end == nothing 

		# first make sure that have initial guess of exp VF 
		if guess_exp_VF == nothing 
			guess_exp_VF = (1/(1-beta))*profits_NC_baseline
		end 

		# next, solve for VF end 
		VF_results_end = get_VF_cf(
			w = w_path[length_transition], Y = Y_path[length_transition],
			param = param, guess_exp_VF = guess_exp_VF)
	end 

	# compute profits end if not provided 
	if profits_object_end == nothing 
		profits_object_end = compute_profits_NC(w = w_path[length_transition], Y = Y_path[length_transition], param = param, z_grid = z_grid) 
	end 

	# Initialize VF path to save results 
	VF_path_end = (VF = VF_results_end.VF, exit_proba = VF_results_end.exit_proba, exp_VF = VF_results_end.exp_VF, 
				   exp_fcost = VF_results_end.exp_fcost, within_period_objects = profits_object_end) 


	### Main loop: Construct value function recursively 

	# create objects needed for iteration 
	VF_path = repeat([VF_path_end],length_transition)
	next_VF_object = deepcopy(VF_path_end)

	for year_from_end in 1:1:(length_transition - 1)

		if verbose
			println("Computing VF for period: ", length_transition - year_from_end, ".")
		end 

		# get new VF object 
		VF_path[length_transition - year_from_end] = get_yearly_VF(
			w = w_path[length_transition - year_from_end], 
			Y = Y_path[length_transition - year_from_end], 
			next_VF_object = next_VF_object)

		# update next_VF_object 
		next_VF_object = deepcopy(VF_path[length_transition - year_from_end])
	end 

	# return results 
	return VF_path
end

# Specify here the type of every input
function get_VF_baseline_DRS(;param,guess_exp_VF,exp_profits,verbose=false,crit = 1e-6,max_iter = 500)
	
	# initialize objects 
	exp_VF_new = copy(guess_exp_VF) 
	exp_VF_old = copy(guess_exp_VF)

	iter = 1
	VF_diff = Inf 
	while VF_diff > crit && iter < max_iter

		# print progress
		if verbose
			println("Iteration ",iter," with difference: ",VF_diff)
		end 
		
		# get exit probability given new guess 
		exit_proba = 1.0 .- exp.(.-exp.(.-(beta.*exp_VF_new .- param.level_fcost)./param.scale_fcost))
		exp_fcost = compute_exp_fcost(exit_proba = exit_proba, exp_VF = exp_VF_new, param = param)

		# compute new VF (over all z)
		VF_new = exp_profits .+ (1.0 .- exit_proba).*( .-exp_fcost .+ beta.*exp_VF_new )
			
		# save old results 
		exp_VF_old = copy(exp_VF_new)
			
		# compute new expected VF 
		exp_VF_new = DRS_transition_z * VF_new 
			
		# Update critical value + iteration
		VF_diff = maximum(abs.(exp_VF_new - exp_VF_old))
		iter = iter + 1		
	end

	## Check if cost parameters need to be indexed by wage!! 
  
	# update results
	exit_proba = 1.0 .- exp.(.-exp.(.-(beta.*exp_VF_new .- param.level_fcost)./param.scale_fcost))
	exp_fcost = compute_exp_fcost(exit_proba = exit_proba, exp_VF = exp_VF_new, param = param)
	VF = exp_profits .+ (1.0 .- exit_proba).*( .-exp_fcost .+ beta.*exp_VF_new )
  
  	# Return results 
	return (exp_VF = exp_VF_new, exit_proba = exit_proba, VF = VF, exp_fcost = exp_fcost)
end

# Specify here the type of every input
function get_VF_cf_DRS(;w,Y,type = "NC",tax=nothing,subsidy=nothing,param,guess_exp_VF,verbose=false,crit = 1e-6,max_iter = 500, n_eps = 500)
	
	# initialize objects 
	exp_VF_new = copy(guess_exp_VF)
	exp_VF_old = copy(guess_exp_VF)

	if tax == nothing
        tax = copy(vat_tax) 
    end 

	# compute z grid based on baseline results & choice of sigma 
	x_bar_baseline = (DRS_baseline_aggregates_goods.total_output)^(1/param.sigma)
	z_grid = (DRS_z_star_grid ./ x_bar_baseline).^(param.sigma/(param.sigma-1))

	# compute z_grid as function of sigma (to be able to vary in counterfactuals) 

	# compute profits
	if type == "NC"
		profit_results = compute_profits_NC(w = w, Y = Y, tax = tax, param = param, z_grid = z_grid)  
		# save results objects 
		exp_profits = profit_results.profits 
		revenue = profit_results.revenue 
		output_aggr = profit_results.output_aggr

	elseif type == "C" 
		profit_results = compute_profits_grid_cf(w = w, Y = Y, tax = tax, n_eps = n_eps, param = param, z_grid = z_grid)
		# save results objects 
		exp_profits = profit_results.expected_profits
		revenue = profit_results.expected_revenue
		output_aggr = profit_results.expected_revenue_output

	elseif type == "C-fix"
		if subsidy == nothing 
			println("Error: Need to specify subsidy rate")
		end 
		profit_results = compute_profits_grid_cf_fix_subsidy(w = w, Y = Y, subsidy = subsidy, tax = nothing, param = param, z_grid = z_grid)
		# save results objects 
		exp_profits = profit_results.expected_profits
		revenue = profit_results.expected_revenue
		output_aggr = profit_results.expected_revenue_output

	else 
		println("No other options than type == {C,NC}")
	end

	iter = 1
	VF_diff = Inf 
	while VF_diff > crit && iter < max_iter

		# print progress
		if verbose
			println("Iteration ",iter," with difference: ",VF_diff)
		end 
		
		# get exit probability given new guess 
		exit_proba = 1.0 .- exp.(.-exp.(.-(beta.*exp_VF_new .- param.level_fcost)./param.scale_fcost))
		exp_fcost = compute_exp_fcost(exit_proba = exit_proba, exp_VF = exp_VF_new, param = param)

		# compute new VF (over all z)
		VF_new = exp_profits .+ (1.0 .- exit_proba).*( .-exp_fcost .+ beta.*exp_VF_new )
			
		# save old results 
		exp_VF_old = copy(exp_VF_new)
			
		# compute new expected VF 
		exp_VF_new = DRS_transition_z * VF_new 
			
		# Update critical value + iteration
		VF_diff = maximum(abs.(exp_VF_new - exp_VF_old))
		iter = iter + 1		
	end

	## Check if cost parameters need to be indexed by wage!! 
  
	# update results
	exit_proba = 1.0 .- exp.(.-exp.(.-(beta.*exp_VF_new .- param.level_fcost)./param.scale_fcost))
	exp_fcost = compute_exp_fcost(exit_proba = exit_proba, exp_VF = exp_VF_new, param = param)
	VF = exp_profits .+ (1.0 .- exit_proba).*( .-exp_fcost .+ beta.*exp_VF_new )
  
  	# Return results 
	return (exp_VF = exp_VF_new, exit_proba = exit_proba, VF = VF, exp_fcost = exp_fcost, 
			exp_profits = exp_profits, revenue = revenue, output_aggr = output_aggr)
end



